import os
import asyncio
from llama_index.core import SimpleDirectoryReader, Settings
from llama_index.llms.deepseek import DeepSeek
from llama_index.core.embeddings import resolve_embed_model
from llama_index.core.response_synthesizers import Refine

# --- 1. 配置区域 ---
# !! 重要：请根据您的环境修改以下路径和密钥 !!

# 设置您的 DeepSeek API 密钥
# 更好的做法是从环境变量加载，以避免将密钥硬编码在代码中
DEEPSEEK_API_KEY = 'sk-f5f3487896554276aaf657fe9b9661a7'

# 设置要读取的文档所在的文件夹路径
INPUT_DIR = r'D:\pythonProject17\refine\f44e1171-2b80-441c-b2cf-d7970d767496\119b3497-6a-3d-46f0-8446-3c5584f1a3a3'

# 设置本地嵌入模型的路径
EMBED_MODEL_PATH = "local:D:/pythonProject17/transformers/model_em/BAAI/bge-small-en-v1.5"

# 设置输出文件的名称
OUTPUT_FILENAME = "document_summaries.txt"


# --- 2. 核心逻辑封装在异步函数中 ---
async def main():
    """
    主执行函数，加载文档、生成摘要并保存到文件。
    """
    # 设置 API 密钥环境变量
    os.environ['DEEPSEEK_API_KEY'] = DEEPSEEK_API_KEY

    # --- LlamaIndex 全局设置 ---
    print("正在配置 LlamaIndex 设置...")
    llm = DeepSeek(model='deepseek-chat')
    embed_model = resolve_embed_model(EMBED_MODEL_PATH)
    Settings.llm = llm
    Settings.embed_model = embed_model
    print("配置完成。")

    # --- 加载文档 ---
    print(f"正在从目录加载文档: {INPUT_DIR}")
    try:
        reader = SimpleDirectoryReader(input_dir=INPUT_DIR)
        docs = reader.load_data()
        if not docs:
            print(f"错误：在目录 '{INPUT_DIR}' 中未找到任何文档。请检查路径。")
            return
        print(f"成功加载 {len(docs)} 个文档。")
    except Exception as e:
        print(f"加载文档时出错: {e}")
        return

    # --- 生成摘要 ---
    summarizer = Refine(llm=llm, verbose=True)
    prompt = "你只要说这个文档做了什么"
    all_summaries = []

    print(f"开始为 {len(docs)} 个文档生成摘要...")

    for i, doc in enumerate(docs):
        print(f"\n正在处理文档 {i + 1}/{len(docs)} (ID: {doc.id_})...")
        document_text = doc.text

        # 调用异步方法 aget_response
        response = await summarizer.aget_response(prompt, [document_text])

        # 将摘要（确保是字符串类型）存入列表
        all_summaries.append(str(response))

        print(f"  -> 摘要生成成功: '{str(response)[:80].replace(os.linesep, ' ')}...'")

    # --- 将所有摘要写入同一个文件 ---
    try:
        with open(OUTPUT_FILENAME, 'w', encoding='utf-8') as f:
            print(f"\n正在将所有摘要写入文件: '{OUTPUT_FILENAME}'...")

            for i, summary in enumerate(all_summaries):
                f.write(f"--- 文档 {i + 1} (ID: {docs[i].id_}) 的摘要 ---\n")
                f.write(summary)
                f.write("\n\n")

        print(f"成功！所有摘要已保存到 {os.path.abspath(OUTPUT_FILENAME)}")

    except IOError as e:
        print(f"错误：写入文件失败。原因: {e}")


# --- 3. 脚本的标准入口点 ---
if __name__ == "__main__":
    # 使用 asyncio.run() 来启动并运行 main 异步函数
    # 这是在 .py 脚本中执行异步代码的标准方式
    asyncio.run(main())